# models/factory.py（极简版，以后新增模型只加elif）
from .lenet import LeNet5  # 从lenet.py导入LeNet5类
from .simple_cnn import SimpleCNN  # 从simple_cnn.py导入SimpleCNN类
from .multi_scale_cnn import MultiScaleCNN  # 从multi_scale_cnn.py导入MultiScaleCNN类

def create_model(model_name, num_classes=10, **kwargs):
    """
    固定规则：model_name是「文件小写名」，对应「文件名.py」中的类
    例：model_name="lenet5" → 对应 lenet.py 中的 LeNet5 类
    """
    if model_name == "lenet5":  # 工厂注册名：文件小写名（lenet）+ 类后缀（5）= lenet5
        return LeNet5(num_classes=num_classes)
    elif model_name == "simple_cnn":  # 工厂注册名：文件小写名（simple_cnn.py → simple_cnn）
        return SimpleCNN(num_classes=num_classes)
    elif model_name == "multi_scale_cnn":  # 工厂注册名：文件小写名（multi_scale_cnn.py → multi_scale_cnn）
        return MultiScaleCNN(num_classes=num_classes, **kwargs)  # 多尺度传参数
    else:
        # 错误提示直接列所有支持的模型名，方便排查
        raise ValueError(
            f"未知模型名称: {model_name}！\n"
            f"支持的模型名：lenet5（对应lenet.py）、simple_cnn（对应simple_cnn.py）、multi_scale_cnn（对应multi_scale_cnn.py）"
        )